refactor: Migrate to `usize` indexing #4273

wingertge · 2026-01-04T20:44:16Z

Pull Request Template

Checklist

Confirmed that cargo run-checks command has been executed.
Made sure the book is up to date with changes in this PR.

Related Issues/PRs

Migrates to changes in tracel-ai/cubecl#1127 and tracel-ai/cubek#51

Changes

Migrates cubecl kernels and fusion to usize indexing

Testing

The test suite runs successfully, though 64-bit indexing is not yet enabled.

laggui

Changes in burn LGTM! Pending rev updates once linked PRs are merged.

laggui · 2026-01-05T13:41:36Z

crates/burn-backend-tests/tests/common/autodiff.rs

    burn_tensor::bf16,
    "../autodiff/mod.rs",
-    ["vulkan", "metal"] // ["cuda", "rocm"] TODO
+    ["metal"] // ["cuda", "rocm"] TODO, ["vulkan"] only supports bf16 for matmul


Whoops, I added bf16 for vulkan but clearly the tests don't pass 😅 thanks for fixing.

[Unrelated to this PR]

I vaguely remember adding it when refactoring the tests since it was a supported global type for vulkan in cubecl. Also reflected by B::supports_dtype:

burn/crates/burn-wgpu/src/lib.rs

Lines 117 to 145 in 8682709

#[test]

fn should_support_dtypes() {

type B = Wgpu;

let device = Default::default();

assert!(B::supports_dtype(&device, DType::F32));

assert!(B::supports_dtype(&device, DType::I64));

assert!(B::supports_dtype(&device, DType::I32));

assert!(B::supports_dtype(&device, DType::U64));

assert!(B::supports_dtype(&device, DType::U32));

assert!(B::supports_dtype(

&device,

DType::QFloat(CubeTensor::<WgpuRuntime>::default_scheme())

));

// Registered as supported type but we don't actually use it?

assert!(B::supports_dtype(&device, DType::Bool));

#[cfg(feature = "vulkan")]

{

assert!(B::supports_dtype(&device, DType::F16));

assert!(B::supports_dtype(&device, DType::BF16));

assert!(B::supports_dtype(&device, DType::I16));

assert!(B::supports_dtype(&device, DType::I8));

assert!(B::supports_dtype(&device, DType::U16));

assert!(B::supports_dtype(&device, DType::U8));

assert!(!B::supports_dtype(&device, DType::F64));

assert!(!B::supports_dtype(&device, DType::Flex32));

}

maybe we should have a better way to represent the actual supported types?

burn/crates/burn-cubecl/src/backend.rs

Lines 83 to 88 in 8682709

fn supports_dtype(device: &Self::Device, dtype: DType) -> bool {

let client = R::client(device);

let ty: StorageType = dtype.into();

client.properties().supports_type(ty.elem_type())

}

That way, tested dtypes can actually reflected supported dtypes.

The supports_type method on the CubeCL side only checks if a type is supported in any way (in this case, it's supported for conversion, as a type for buffers, and for dot product on Intel, along with tensor core instructions). It's kinda tough though because there's no good way to express that in just a single boolean (hence why the TypeUsage enum set exists in CubeCL).

This is how it's registered for Vulkan

if let Some(bfloat16) = ext_feat.bfloat16 { if bfloat16.shader_b_float16_type == TRUE { register( ElemType::Float(FloatKind::BF16).into(), TypeUsage::Conversion | TypeUsage::Buffer, ); } if bfloat16.shader_b_float16_dot_product == TRUE { register( ElemType::Float(FloatKind::BF16).into(), TypeUsage::DotProduct.into(), ); } }

So it's supported for matmul and casting, but none of the other ops.

The supports_type method on the CubeCL side only checks if a type is supported in any way

Yeah and for the first draft I simply mirrored that for the backends, but I think it should be refined.

It's kinda tough though because there's no good way to express that in just a single boolean (hence why the TypeUsage enum set exists in CubeCL).

That's a good point. It's still useful to query backend supported types for burn, so maybe we should also define an enum similar to TypeUsage? (without atomics, and perhaps consolidate conversion / buffer into "storage" variant or similar).

Regardless, that's something for a separate issue/PR I think.

Oh yes entirely, that's why I prefaced my initial comment with [Unrelated to this PR] 😄

Just wanted to get your thoughts since this was related to the bf16 change

codecov · 2026-01-08T17:36:52Z

Codecov Report

❌ Patch coverage is 0% with 831 lines in your changes missing coverage. Please review.
✅ Project coverage is 68.87%. Comparing base (d682723) to head (e0d9252).
⚠️ Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/burn-cubecl-fusion/src/engine/codegen/io.rs	0.00%	89 Missing ⚠️
.../cube/connected_components/hardware_accelerated.rs	0.00%	77 Missing ⚠️
...es/burn-cubecl-fusion/src/engine/codegen/kernel.rs	0.00%	54 Missing ⚠️
crates/burn-cubecl-fusion/src/optim/matmul/args.rs	0.00%	39 Missing ⚠️
...ecl-fusion/src/engine/launch/vectorization/base.rs	0.00%	35 Missing ⚠️
...c/backends/cube/connected_components/prefix_sum.rs	0.00%	34 Missing ⚠️
crates/burn-cubecl/src/kernel/conv/direct.rs	0.00%	28 Missing ⚠️
...rates/burn-cubecl/src/kernel/index/slice_assign.rs	0.00%	27 Missing ⚠️
...ates/burn-cubecl-fusion/src/engine/codegen/view.rs	0.00%	26 Missing ⚠️
...-cubecl/src/kernel/conv/deform_conv_transpose2d.rs	0.00%	26 Missing ⚠️
... and 53 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4273      +/-   ##
==========================================
+ Coverage   68.85%   68.87%   +0.02%     
==========================================
  Files        1405     1405              
  Lines      167686   167607      -79     
==========================================
- Hits       115456   115442      -14     
+ Misses      52230    52165      -65

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

wingertge added 3 commits January 4, 2026 21:40

Migrate to usize indexing

c5f0a50

Merge branch 'main' into feat/usize

91bc2d3

Fix merge

e60a2d3

laggui reviewed Jan 5, 2026

View reviewed changes

wingertge added 3 commits January 7, 2026 23:42

Update cubecl

6f2c8d4

Merge branch 'main' into feat/usize

6fcf199

Update cubek

e0d9252

wingertge requested a review from laggui January 8, 2026 16:46

laggui approved these changes Jan 8, 2026

View reviewed changes

laggui merged commit 7b32614 into tracel-ai:main Jan 8, 2026
10 checks passed

wingertge deleted the feat/usize branch January 8, 2026 21:16

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

refactor: Migrate to `usize` indexing #4273

refactor: Migrate to `usize` indexing #4273

Uh oh!

wingertge commented Jan 4, 2026

Uh oh!

laggui left a comment

Uh oh!

laggui Jan 5, 2026

Uh oh!

wingertge Jan 5, 2026 •

edited

Loading

Uh oh!

wingertge Jan 5, 2026

Uh oh!

laggui Jan 5, 2026

Uh oh!

wingertge Jan 8, 2026

Uh oh!

laggui Jan 8, 2026 •

edited

Loading

Uh oh!

codecov bot commented Jan 8, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	#[test]
	fn should_support_dtypes() {
	type B = Wgpu;
	let device = Default::default();

	assert!(B::supports_dtype(&device, DType::F32));
	assert!(B::supports_dtype(&device, DType::I64));
	assert!(B::supports_dtype(&device, DType::I32));
	assert!(B::supports_dtype(&device, DType::U64));
	assert!(B::supports_dtype(&device, DType::U32));
	assert!(B::supports_dtype(
	&device,
	DType::QFloat(CubeTensor::<WgpuRuntime>::default_scheme())
	));
	// Registered as supported type but we don't actually use it?
	assert!(B::supports_dtype(&device, DType::Bool));

	#[cfg(feature = "vulkan")]
	{
	assert!(B::supports_dtype(&device, DType::F16));
	assert!(B::supports_dtype(&device, DType::BF16));
	assert!(B::supports_dtype(&device, DType::I16));
	assert!(B::supports_dtype(&device, DType::I8));
	assert!(B::supports_dtype(&device, DType::U16));
	assert!(B::supports_dtype(&device, DType::U8));

	assert!(!B::supports_dtype(&device, DType::F64));
	assert!(!B::supports_dtype(&device, DType::Flex32));
	}

	fn supports_dtype(device: &Self::Device, dtype: DType) -> bool {
	let client = R::client(device);

	let ty: StorageType = dtype.into();
	client.properties().supports_type(ty.elem_type())
	}

refactor: Migrate to usize indexing #4273

refactor: Migrate to usize indexing #4273

Uh oh!

Conversation

wingertge commented Jan 4, 2026

Pull Request Template

Checklist

Related Issues/PRs

Changes

Testing

Uh oh!

laggui left a comment

Choose a reason for hiding this comment

Uh oh!

laggui Jan 5, 2026

Choose a reason for hiding this comment

Uh oh!

wingertge Jan 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wingertge Jan 5, 2026

Choose a reason for hiding this comment

Uh oh!

laggui Jan 5, 2026

Choose a reason for hiding this comment

Uh oh!

wingertge Jan 8, 2026

Choose a reason for hiding this comment

Uh oh!

laggui Jan 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov bot commented Jan 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

refactor: Migrate to `usize` indexing #4273

refactor: Migrate to `usize` indexing #4273

wingertge Jan 5, 2026 •

edited

Loading

laggui Jan 8, 2026 •

edited

Loading

codecov bot commented Jan 8, 2026 •

edited

Loading